import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from xgboost import XGBRegressor
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import mean_squared_error, r2_score

# ✅ 设置中文字体（避免图中乱码）
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False
sns.set(style='whitegrid')

# 1. 数据读取
df = pd.read_csv('../data/cleaned_data.csv', encoding='gbk')

# 2. 时间戳 + 衍生特征
df['shijian'] = pd.to_datetime(df['date'].astype(str) + ' ' + df['time'].astype(str), errors='coerce')
df['hour'] = df['shijian'].dt.hour
df['day'] = df['shijian'].dt.dayofweek
df['KH_IM_ratio'] = df['shuliaoKH'] / (df['shuliaoIM'] + 1e-6)

# 特征列表构建
features = [
    'weiliaoc', 'yaotouc', 'bilengjydS1', 'bilengjedS1',
    'bilengjsdI1', 'shuliaoKH', 'shuliaoIM',
    'fengjizs', 'rehao', 'hour', 'day', 'KH_IM_ratio'
]

# 检查其他可能字段
for col in ['shuliaol', 'chumoslKH', 'chumoslSM']:
    if col in df.columns:
        features.append(col)

target = 'shuliaoSM'

# 3. 编码类别字段 & 缺失处理
df = df[features + [target]].dropna()
for col in ['weiliaoc', 'yaotouc']:
    if df[col].dtype == 'object':
        df[col] = LabelEncoder().fit_transform(df[col])

X = df[features]
y = df[target]

# 4. 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 5. XGBoost 模型训练
model = XGBRegressor(
    n_estimators=200,
    max_depth=6,
    learning_rate=0.1,
    subsample=0.8,
    colsample_bytree=0.8,
    random_state=42,
    verbosity=1
)
model.fit(X_train, y_train)

# 6. 预测与评估
y_pred = model.predict(X_test)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
r2 = r2_score(y_test, y_pred)

print("✅ XGBoost 模型训练完成")
print(f"RMSE: {rmse:.4f}")
print(f"R²: {r2:.4f}")

# 7. 保存预测结果
output_df = pd.DataFrame({
    'True_SM': y_test.values,
    'Predicted_SM': y_pred
})
output_df['Error'] = output_df['Predicted_SM'] - output_df['True_SM']
output_df.to_csv('../data/xgb_prediction_results.csv', index=False, encoding='gbk')
print("📁 预测结果已保存至 xgb_prediction_results.csv")

# 8. 可视化（预测效果 + 残差分析）
plt.figure(figsize=(6,6))
plt.scatter(output_df['True_SM'], output_df['Predicted_SM'], alpha=0.6, edgecolors='k')
plt.plot([y.min(), y.max()], [y.min(), y.max()], '--', color='gray')
plt.xlabel("实际饱和比")
plt.ylabel("预测饱和比")
plt.title("① 预测值 vs 实际值")
plt.grid(True)
plt.tight_layout()
plt.show()

plt.figure(figsize=(8,4))
sns.histplot(output_df['Error'], bins=30, kde=True, color='orange')
plt.axvline(0, linestyle='--', color='gray')
plt.title("② 预测误差分布图")
plt.xlabel("预测误差")
plt.tight_layout()
plt.show()

plt.figure(figsize=(6,4))
sns.boxplot(x=output_df['Error'], color='skyblue')
plt.title("③ 预测误差箱线图")
plt.tight_layout()
plt.show()

plt.figure(figsize=(6,4))
plt.scatter(output_df['Predicted_SM'], output_df['Error'], alpha=0.5, edgecolors='k')
plt.axhline(0, linestyle='--', color='gray')
plt.xlabel("预测饱和比")
plt.ylabel("残差（预测 - 实际）")
plt.title("④ 残差 vs 预测值")
plt.tight_layout()
plt.show()
